import { create } from "zustand";
import { persist } from "zustand/middleware";
import { useAuthStore } from "./auth";
import { requestStream } from "@/api/stream";
import { Toast, Dialog } from "antd-mobile";
import {
  requestWithPrompt,
  addConversation,
  deleteConversation,
  getConversationHistory,
  IConversation,
  updateConversationName,
  moveConversation,
  collectConver,
  cancelCollectConver,
  deleteNotCollectConver,
} from "@/api/conversation";
import { updateAvatarApi } from "@/api/auth";
import { trimTopic, isUrlExceedLimit, getVersion } from "@/utils/utils";
import { LOCAL_KEY } from "@/constant/env";
import Locale from "../locales";
import dayjs from "dayjs";

export const APP_KEY = LOCAL_KEY + "__app";

export type Message = {
  role: string;
  content: string;
  date: string;
  streaming?: boolean;
};

export enum SubmitKey {
  Enter = "Enter",
  CtrlEnter = "Ctrl + Enter",
  ShiftEnter = "Shift + Enter",
  AltEnter = "Alt + Enter",
  MetaEnter = "Meta + Enter",
}

export enum Theme {
  // Auto = "auto",
  Dark = "dark",
  Light = "light",
}

export interface ChatConfig {
  historyMessageCount: number; // -1 means all
  compressMessageLengthThreshold: number;
  sendBotMessages: boolean; // send bot's message or not
  submitKey: SubmitKey;
  avatar: string;
  fontSize: number;
  theme: Theme;
  tightBorder: boolean;
  disablePromptHint: boolean;
}

export function isValidNumber(x: number, min: number, max: number) {
  return typeof x === "number" && x <= max && x >= min;
}

const DEFAULT_CONFIG: ChatConfig = {
  historyMessageCount: 4,
  compressMessageLengthThreshold: 1000,
  sendBotMessages: true as boolean,
  submitKey: SubmitKey.CtrlEnter as SubmitKey,
  avatar: "1f603",
  fontSize: 14,
  theme: Theme.Light as Theme,
  tightBorder: true,
  disablePromptHint: false,
};
export interface ChatStat {
  wordCount: number;
  charCount: number;
}
export interface ChatSession {
  id: string;
  topic: string;
  memoryPrompt: string;
  messages: Message[];
  isCollect: boolean;
  stat: ChatStat;
  lastUpdate: string;
  lastSummarizeIndex: number;
}

const DEFAULT_TOPIC = Locale.Store.DefaultTopic;

function createEmptySession(conversation?: IConversation): ChatSession {
  const createDate = dayjs().format("YYYY/MM/DD HH:mm:ss");
  const result = {
    id: "",
    topic: DEFAULT_TOPIC,
    memoryPrompt: "",
    isCollect: false,
    messages: [
      {
        role: "assistant",
        content: Locale.Store.BotHello,
        date: createDate,
      },
    ],
    stat: {
      wordCount: 0,
      charCount: 0,
    },
    lastUpdate: createDate,
    lastSummarizeIndex: 0,
  };
  if (conversation) {
    Object.assign(result, {
      ...result,
      id: conversation.id,
      topic: conversation.conversationName,
      lastUpdate: conversation.createTime,
      isCollect: conversation.isCollect === "0",
    });
    result.stat.charCount = conversation.charCount;
  }
  return result;
}

interface ChatStore {
  isLoading: boolean;
  sessions: ChatSession[];
  config: ChatConfig;
  currentSessionIndex: number;
  autoScroll: boolean;
  setLoading: (isLoading: boolean) => void;
  setAutoScroll: (autoScroll: boolean) => void;
  initialSession: (conversations: IConversation[]) => void;
  updateSession: (session: ChatSession[]) => void;
  updateSessionIndex: (index: number) => void;
  removeSession: (index: number) => Promise<void>;
  selectSession: (index: number) => void;
  getSessionHistory: (conversationId: string) => Promise<Message[]>;
  newSession: () => void;
  currentSession: () => ChatSession;
  onNewMessage: (message: Message) => void;
  onUserInput: (content: string) => Promise<void>;
  summarizeSession: (
    content: string,
    session: ChatSession,
    index: number,
  ) => void;
  updateStat: (message: Message) => void;
  updateCurrentSession: (updater: (session: ChatSession) => void) => void;
  updateCurrentSessionName: (id: string, topic: string) => Promise<void>;
  updateMessage: (
    sessionIndex: number,
    messageIndex: number,
    updater: (message?: Message) => void,
  ) => void;
  getMessagesWithMemory: () => Message[];
  getMemoryPrompt: () => Message;
  getConfig: () => ChatConfig;
  resetConfig: () => void;
  toggleCollectSession: (session: ChatSession) => Promise<void>;
  updateConfig: (updater: (config: ChatConfig) => void) => void;
  updateAvatar: (avatar: string) => Promise<void>;
  updateConversationIndex: (
    conversationId: string,
    downConversionId: string | null,
  ) => Promise<void>;
  clearAllData: () => Promise<void>;
}

function countMessages(msgs: Message[]) {
  return msgs.reduce((pre, cur) => pre + cur.content.length, 0);
}
export const useChatStore = create(
  persist<ChatStore>(
    (set, get) => ({
      isLoading: false,
      sessions: [createEmptySession()],
      currentSessionIndex: 0,
      config: {
        ...DEFAULT_CONFIG,
      },
      autoScroll: false,
      setLoading(isLoading) {
        set({ isLoading });
      },
      setAutoScroll(autoScroll) {
        set({ autoScroll });
      },
      initialSession(conversations) {
        const sessions: ChatSession[] = [];
        conversations.map((v) => {
          sessions.push(createEmptySession(v));
        });
        set({
          sessions: [createEmptySession()].concat(sessions),
        });
      },
      updateSession(sessions) {
        set({
          sessions,
        });
      },
      async removeSession(index) {
        let nextIndex = get().currentSessionIndex;
        const sessions = get().sessions;
        let conversationId = sessions[index].id;
        const cb = (): void => {
          set((state) => {
            let nextIndex = state.currentSessionIndex;
            const sessions = state.sessions;

            if (sessions.length === 1) {
              return {
                currentSessionIndex: 0,
                sessions: [createEmptySession()],
              };
            }

            sessions.splice(index, 1);

            if (nextIndex === index) {
              nextIndex -= 1;
            }
            Toast.show({
              icon: "success",
              content: "删除成功",
            });
            return {
              currentSessionIndex: nextIndex,
              sessions,
            };
          });

          // sessions.splice(index, 1);
          // if (nextIndex === index) {
          //   nextIndex -= 1;
          // }
          // set({
          //   currentSessionIndex: nextIndex,
          //   sessions,
          // });
        };
        if (conversationId) {
          const authStore = useAuthStore.getState();
          // if (sessions.length === 1) return;
          await deleteConversation({
            conversationId,
            userId: authStore.user.id,
          });
          cb();
        } else {
          cb();
        }
      },
      selectSession(index) {
        set({
          currentSessionIndex: index,
        });
        const sessions = get().sessions;
        const session = sessions[index];
        if (!session.id) return;
        get()
          .getSessionHistory(session.id)
          .then((messages: Message[]) => {
            sessions[index].messages = messages;
            set({ sessions });
          });
      },
      updateSessionIndex(index) {
        set({
          currentSessionIndex: index,
        });
      },
      async getSessionHistory(conversationId) {
        const authStore = useAuthStore.getState();
        const response = await getConversationHistory({
          conversationId,
          userId: authStore.user.id,
        });
        return response;
      },
      newSession() {
        set((state) => ({
          sessions: [createEmptySession()].concat(state.sessions),
        }));
        set((state) => ({
          currentSessionIndex: 0,
        }));
        Toast.show({
          icon: "success",
          content: "新的聊天",
        });
        // addConversation({
        //   userId: authStore.user.id,
        //   appUser: authStore.user.appUser,
        // }).then((response) => {
        //   const session = {
        //     ...createEmptySession(),
        //     id: response.id,
        //   };
        //   set((state) => ({
        //     sessions: [session].concat(state.sessions),
        //   }));
        //   set((state) => ({
        //     currentSessionIndex: 0,
        //   }));
        //   Toast.show({
        //     icon: 'success',
        //     content: '新的聊天',
        //   })
        // });
      },
      currentSession() {
        let index = get().currentSessionIndex;
        const sessions = get().sessions;
        if (index < 0 || index >= sessions.length) {
          index = Math.min(sessions.length - 1, Math.max(0, index));
          set(() => ({ currentSessionIndex: index }));
        }
        const session = sessions[index];
        return session;
      },
      onNewMessage(message) {
        get().updateCurrentSession((session) => {
          session.lastUpdate = dayjs().format("YYYY/MM/DD HH:mm:ss");
        });
        get().updateStat(message);
      },
      async onUserInput(content) {
        const authStore = useAuthStore.getState();
        if (isUrlExceedLimit(content)) {
          Toast.show({
            icon: "fail",
            content: "输入文字长度超出限制",
          });
          return;
        }
        const userMessage: Message = {
          role: "user",
          content,
          date: dayjs().format("YYYY/MM/DD HH:mm:ss"),
        };
        const botMessage: Message = {
          content: "",
          role: "assistant",
          date: dayjs().format("YYYY/MM/DD HH:mm:ss"),
          streaming: true,
        };
        let conversationId = get().currentSession().id;
        const currentSession = get().currentSession();
        const currentSessionIndex = get().currentSessionIndex;
        if (!conversationId) {
          const response = await addConversation({
            userId: authStore.user.id,
            appUser: authStore.user.appUser,
          });
          conversationId = response.id;
          get().updateCurrentSession((session) => {
            session.id = conversationId;
          });
        }

        // save user's and bot's message
        get().updateCurrentSession((session) => {
          session.messages.push(userMessage);
          session.messages.push(botMessage);
        });
        // 获取topic
        get().summarizeSession(content, currentSession, currentSessionIndex);
        // make request
        console.log("[User Input] ", userMessage);
        get().setAutoScroll(true);
        requestStream(conversationId, content, {
          message: (content, done) => {
            if (done) {
              botMessage.streaming = false;
              botMessage.content = content;
              get().onNewMessage(botMessage);
              get().setAutoScroll(false);
            } else {
              botMessage.content = content;
            }
            set({});
          },
          error: () => {
            botMessage.streaming = false;
            get().setAutoScroll(false);
            botMessage.content += "\n\n" + Locale.Store.Error;
            set({});
          },
        });
      },
      summarizeSession(prompt, session, index) {
        if (session.topic === DEFAULT_TOPIC || !session.topic) {
          requestWithPrompt(session.id, prompt).then((res) => {
            console.log(
              "🚀 ~ file: app.ts:396 ~ requestWithPrompt ~ res:",
              res,
            ); 
            set((state) => ({
              sessions: state.sessions.map((s, i) =>
                i === index ? { ...s, topic: res } : s,
              ),
            }));
          });
        }

        // should summarize topic after chating more than 50 words
        // const SUMMARIZE_MIN_LEN = 50;
        // if (
        //   session.topic === DEFAULT_TOPIC &&
        //   countMessages(session.messages) >= SUMMARIZE_MIN_LEN
        // ) {
        //   // TODO 传递会话的id和Topic
        //   requestWithPrompt(session.id, Locale.Store.Prompt.Topic).then(
        //     (res) => {
        //       get().updateCurrentSession(
        //         (session) => (session.topic = trimTopic(res)),
        //       );
        //     },
        //   );
        // }

        // const config = get().config;
        // let toBeSummarizedMsgs = session.messages.slice(
        //   session.lastSummarizeIndex,
        // );
        // const historyMsgLength = countMessages(toBeSummarizedMsgs);

        // if (historyMsgLength > 4000) {
        //   toBeSummarizedMsgs = toBeSummarizedMsgs.slice(
        //     -config.historyMessageCount,
        //   );
        // }

        // // add memory prompt
        // toBeSummarizedMsgs.unshift(get().getMemoryPrompt());

        // const lastSummarizeIndex = session.messages.length;

        // console.log(
        //   "[Chat History] ",
        //   toBeSummarizedMsgs,
        //   historyMsgLength,
        //   config.compressMessageLengthThreshold,
        // );

        // if (historyMsgLength > config.compressMessageLengthThreshold) {
        //   requestWithPrompt(session.id, Locale.Store.Prompt.Summarize).then(
        //     (res) => {
        //       get().updateCurrentSession(
        //         (session) => (session.topic = trimTopic(res)),
        //       );
        //       session.lastSummarizeIndex = lastSummarizeIndex;
        //     },
        //   );
        // }
      },
      updateMessage(sessionIndex: number, messageIndex: number, updater) {
        const sessions = get().sessions;
        const session = sessions.at(sessionIndex);
        const messages = session?.messages;
        updater(messages?.at(messageIndex));
        set(() => ({ sessions }));
      },

      updateStat(message) {
        get().updateCurrentSession((session) => {
          session.stat.charCount += 1;
          // TODO: should update chat count and word count
        });
      },
      async updateCurrentSessionName(id, topic) {
        await updateConversationName({
          conversationId: id,
          conversationName: topic,
        });
        get().updateCurrentSession((session) => (session.topic = topic!));
      },
      updateCurrentSession(updater) {
        const sessions = get().sessions;
        const index = get().currentSessionIndex;
        updater(sessions[index]);
        set(() => ({ sessions }));
      },
      getMemoryPrompt() {
        const session = get().currentSession();

        return {
          role: "system",
          content: Locale.Store.Prompt.History(session.memoryPrompt),
          date: "",
        } as Message;
      },
      getMessagesWithMemory() {
        const session = get().currentSession();
        const config = get().config;
        const n = session.messages.length;
        const recentMessages = session.messages.slice(
          n - config.historyMessageCount,
        );

        const memoryPrompt = get().getMemoryPrompt();

        if (session.memoryPrompt) {
          recentMessages.unshift(memoryPrompt);
        }
        return recentMessages;
      },
      async resetConfig() {
        const result = await Dialog.confirm({
          content: Locale.Settings.Actions.ResetAll,
        });
        if (result) {
          set(() => ({ config: { ...DEFAULT_CONFIG } }));
        }
      },
      getConfig() {
        return get().config;
      },
      async toggleCollectSession(session) {
        if (session.isCollect) {
          await cancelCollectConver(session.id);
          Toast.show({
            icon: "success",
            content: "取消收藏",
          });
        } else {
          await collectConver(session.id);
          Toast.show({
            icon: "success",
            content: "收藏成功",
          });
        }
        get().updateCurrentSession((session) => {
          session.isCollect = !session.isCollect;
        });
      },
      updateConfig(updater) {
        const config = get().config;
        updater(config);
        set(() => ({ config }));
      },
      async updateConversationIndex(conversationId, downConversionId) {
        await moveConversation({ conversationId, downConversionId });
      },
      async updateAvatar(avatar) {
        const authStore = useAuthStore.getState();
        await updateAvatarApi({
          id: authStore.user.id,
          url: avatar,
        });
        get().updateConfig((config) => (config.avatar = avatar));
      },
      clearAllData: async () => {
        const result = await Dialog.confirm({
          content: Locale.Settings.Actions.ClearAll,
        });
        if (result) {
          const authStore = useAuthStore.getState();
          await deleteNotCollectConver(authStore.user.id);
          // localStorage.clear();
          location.reload();
        }
      },
    }),
    {
      name: APP_KEY,
      version: getVersion(),
    },
  ),
);
